#ifndef	ema_included
#define	ema_included

__global__	void	_emaf(uint32_t	C,	half	*inp,	half	*out,	half	*rnn,	half	*sum) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x;
    half	*p=inp+id*C,	*q=out+id*C,	*s=sum+id*C;
    float	b=1/(1+expf(__half2float(rnn[id])));
    float8	o8= {};
    for(unsigned	c=0;	c<C;	c+=8) {
        float8	i8=_b2f8(*(half8*)(p+c)),s8;
        o8.x.x=(s8.x.x=o8.w.y)+b*(i8.x.x-o8.w.y);
        o8.x.y=(s8.x.y=o8.x.x)+b*(i8.x.y-o8.x.x);
        o8.y.x=(s8.y.x=o8.x.y)+b*(i8.y.x-o8.x.y);
        o8.y.y=(s8.y.y=o8.y.x)+b*(i8.y.y-o8.y.x);
        o8.z.x=(s8.z.x=o8.y.y)+b*(i8.z.x-o8.y.y);
        o8.z.y=(s8.z.y=o8.z.x)+b*(i8.z.y-o8.z.x);
        o8.w.x=(s8.w.x=o8.z.y)+b*(i8.w.x-o8.z.y);
        o8.w.y=(s8.w.y=o8.w.x)+b*(i8.w.y-o8.w.x);
        *(half8*)(q+c)=_f2b8(o8);
        *(half8*)(s+c)=_f2b8(s8);
    }
}

__global__	void	_emab(uint32_t	C,	half	*gin,	half	*gra,	half	*rnn,	half	*rmo,	half	*sum,	half	*inp) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x;
    half	*p=gin+id*C,	*q=gra+id*C,	*s=sum+id*C,	*t=inp+id*C;
    float	b=1/(1+expf(__half2float(rnn[id]))),dr=0;
    float8	o8= {};
    for(unsigned	c=C-1;	c<C;	c-=8) {
        float8	i8=_b2f8(*(half8*)(p+c-7)),	s8=_b2f8(*(half8*)(s+c-7)),	t8=_b2f8(*(half8*)(t+c-7));
        dr+=(s8.w.y-t8.w.y)*(o8.w.y=o8.x.x+b*(i8.w.y-o8.x.x));
        dr+=(s8.w.x-t8.w.x)*(o8.w.x=o8.w.y+b*(i8.w.x-o8.w.y));
        dr+=(s8.z.y-t8.z.y)*(o8.z.y=o8.w.x+b*(i8.z.y-o8.w.x));
        dr+=(s8.z.x-t8.z.x)*(o8.z.x=o8.z.y+b*(i8.z.x-o8.z.y));
        dr+=(s8.y.y-t8.y.y)*(o8.y.y=o8.z.x+b*(i8.y.y-o8.z.x));
        dr+=(s8.y.x-t8.y.x)*(o8.y.x=o8.y.y+b*(i8.y.x-o8.y.y));
        dr+=(s8.x.y-t8.x.y)*(o8.x.y=o8.y.x+b*(i8.x.y-o8.y.x));
        dr+=(s8.x.x-t8.x.x)*(o8.x.x=o8.x.y+b*(i8.x.x-o8.x.y));
        *(half8*)(q+c-7)=_f2b8(o8);
    }
    float	m=__half2float(rmo[id]);
    m+=parameter_beta*(dr*b*(1-b)-m);
    rmo[id]=__float2half_rn(m);
}

template<uint32_t	R,	uint32_t	C>
struct	ema {
    static	tensor<half,R*C>	t0;
    parameter<R>	r;
    tensor<half,C*R>	sum;

    void	save(FILE	*F) {
        r.save(F);
    }

    void	load(FILE	*F) {
        r.load(F);
    }

    uint64_t	size(void) {
        return	r.size();
    }

    void	fw(tensor<half,R*C>	&inp/*const*/,tensor<half,R*C>	&out/*changed*/) {
        r.z2w();
        uint32_t	w=warp_size(1,R);
        _emaf<<<R/w,w>>>(C,inp.p,t0.p,r.w.p,sum.p);
        transpose<C,R>(t0,out);
    }

    void	bk(tensor<half,R*C>	&inp/*const*/,tensor<half,R*C>	&gin/*const*/,tensor<half,R*C>	&gra/*changed*/) {
        transpose<R,C>(gin,gra);
        uint32_t	w=warp_size(1,R);
        _emab<<<R/w,w>>>(C,gra.p,t0.p,r.w.p,r.m.p,sum.p,inp.p);
        transpose<C,R>(t0,gra);
        r.m2z(etalr);
    }
};
template<uint32_t	R,	uint32_t	C>
tensor<half,R*C>	ema<R,C>::t0;
#endif

